#include "ASTInstrumentation.h"
#include "ASTNodeFactory.h"

#include <clang/AST/AST.h>
#include <clang/Sema/Sema.h>

#include <sstream>

void ASTInstrumentation::instrumentTranslationUnit() {
  /// Create an external getenv() declaration within extern "C" {} block.
  clang::LinkageSpecDecl *cLinkageSpecDecl =
      clang::LinkageSpecDecl::Create(context,
                                     context.getTranslationUnitDecl(),
                                     NULL_LOCATION,
                                     NULL_LOCATION,
                                     clang::LinkageSpecDecl::LanguageIDs::lang_c,
                                     true);
  getenvFuncDecl = createGetEnvFuncDecl(cLinkageSpecDecl);
  cLinkageSpecDecl->addDecl(getenvFuncDecl);
  context.getTranslationUnitDecl()->addDecl(cLinkageSpecDecl);
}

void ASTInstrumentation::addMutantStringDefinition(std::string identifier, int mutator, int line,
                                                   int column) {
  std::ostringstream mis;

  mis << "mull_mutation_" << mutator << "_" << line << "_" << column;
  std::string variableName = mis.str();

  clang::IdentifierInfo &varDeclIdentifierInfo = context.Idents.get(variableName);

  clang::StringLiteral *literal = factory.createStringLiteral(identifier);
  literal->setValueKind(clang::VK_RValue);

  clang::QualType qualType =
      factory.getStringLiteralArrayType(context.getConstType(context.CharTy), identifier.size());

  clang::VarDecl *varDecl = clang::VarDecl::Create(context,
                                                   context.getTranslationUnitDecl(),
                                                   NULL_LOCATION,
                                                   NULL_LOCATION,
                                                   &varDeclIdentifierInfo,
                                                   qualType,
                                                   context.getTrivialTypeSourceInfo(qualType),
                                                   clang::StorageClass::SC_Extern);
  varDecl->setInit(literal);

#if defined __APPLE__
  const char *mullSection = "__mull,.mull_mutants";
#else
  const char *mullSection = ".mull_mutants";
#endif

  clang::SectionAttr *mullSectionAttr = factory.createSectionAttr(mullSection);
  varDecl->addAttr(mullSectionAttr);
  context.getTranslationUnitDecl()->addDecl(varDecl);

  assert(varDecl->isThisDeclarationADefinition() == clang::VarDecl::Definition);

  /// This is the most sensitive line of the whole Clang AST-related
  /// instrumentation and mutation done by Mull. Currently it seems to be the
  /// only way to add a new function declaration to an existing Clang AST.
  /// See clang::ParseAST() function with the following line:
  /// """
  ///   // Process any TopLevelDecls generated by #pragma weak.
  ///   for (Decl *D : S.WeakTopLevelDecls())
  ///     Consumer->HandleTopLevelDecl(DeclGroupRef(D));
  ///
  ///   Consumer->HandleTranslationUnit(S.getASTContext());
  /// """
  /// When a declaration is added to the weak top-level decls, the call
  /// HandleTopLevelDecl() with one of the codegen hooks, generates the
  /// actual definitions. Without this workaround, only a declaration of a new
  /// function will appear in the generated LLVM IR.
  /// TODO: Is it the only place to make the codegen generate the function definition?
  sema.WeakTopLevelDecls().push_back(varDecl);
}

clang::FunctionDecl *ASTInstrumentation::getGetenvFuncDecl() {
  assert(getenvFuncDecl);
  return getenvFuncDecl;
}

clang::FunctionDecl *ASTInstrumentation::createGetEnvFuncDecl(clang::DeclContext *declContext) {
  clang::IdentifierInfo &getenvParamNameInfo = context.Idents.get("name");

  clang::FunctionProtoType::ExtProtoInfo ext;

  clang::QualType parameterType = context.getPointerType(context.getConstType(context.CharTy));
  clang::QualType returnType = context.getPointerType(context.CharTy);

  std::vector<clang::QualType> paramTypes;
  paramTypes.push_back(parameterType);

  clang::QualType getenvFuncType = context.getFunctionType(returnType, paramTypes, ext);

  clang::FunctionDecl *getEnvFuncDecl =
      factory.createFunctionDecl("getenv", declContext, getenvFuncType);

  clang::ParmVarDecl *ParDecl =
      clang::ParmVarDecl::Create(context,
                                 getEnvFuncDecl,
                                 NULL_LOCATION,
                                 NULL_LOCATION,
                                 &getenvParamNameInfo,
                                 parameterType,
                                 context.getTrivialTypeSourceInfo(parameterType),
                                 clang::StorageClass::SC_None,
                                 nullptr // Expr *DefArg
      );

  std::vector<clang::ParmVarDecl *> paramDecls;
  paramDecls.push_back(ParDecl);
  getEnvFuncDecl->setParams(paramDecls);

  return getEnvFuncDecl;
}
